import json
import jsonlines
import os


DataSet = [
    "director",
    "emsplus",
    "rcp",
    "umac"
]


def Gen_Ext_Slice_Json(data_name:str,ext_lenth_arg:int=512):

    input_keyword_file_path = f"ext_slice/inputs/output_rag_keyword_{data_name}.json"
    with open(input_keyword_file_path,"r",encoding="utf-8") as file:
        data = json.load(file)
    
    for idx in range(len(data)):
        one = data[idx]

        # id = one["id"]
        # query = one["query"]
        selected_file_path = list(one["selected_file_path"])


        for i in range(len(selected_file_path)):
            res = selected_file_path[i]

            file_path = res["file_path"]
            score = res["score"]
            text_slice = str(res["text_slice"])
            
            try:
                # 读取文件完整内容
                with open(file_path,"r",encoding="utf-8") as file:
                    content = str(file.read())
            except Exception as e:
                print(e)
                continue
            
            # 字段拓展
            ori_len = len(text_slice)
            index = content.index(text_slice)
            back_Ext = ext_lenth_arg + ori_len
            lenth = len(content)
            ext_lenth = min(back_Ext,lenth-index)
            ext_text_slice = content[index:index+ext_lenth]

            # 修改字段
            data[idx]["selected_file_path"][i]["text_slice"] = ext_text_slice
    
    with open(f"ext_slice/outputs/output_rag_ext_{ext_lenth_arg}_slice_{data_name}.json","w",encoding="utf-8") as file:
        json.dump(data,file,ensure_ascii=False,indent=4)


def LoadJsonl(file_path) -> list:
    data = []
    with jsonlines.open(file_path) as reader:
        for obj in reader:
            data.append(obj)
    return data


def Merge_Emsplus_Into_Jsonl(ext_len:int=0):

    emsplus_jsonl_file_path = "/Users/savenneer/工作内容/aiops/datapre/ext_slice/emsplus.jsonl"
    emsplus_data = LoadJsonl(emsplus_jsonl_file_path)

    else_result_jsonl = f"ext_slice/extslice_output/result_{ext_len}.jsonl"
    else_data = LoadJsonl(else_result_jsonl)

    # 合并
    data = emsplus_data + else_data

    data = sorted(data,key=lambda x:x["id"])

    # dump data to jsonl
    with jsonlines.open(f"ext_slice/submit_output/submit_{ext_len}.jsonl","w") as writer:
        for one in data:
            writer.write(one)


if __name__ == "__main__":

    # Merge_Emsplus_Into_Jsonl(0)
    # Merge_Emsplus_Into_Jsonl(256)
    # Merge_Emsplus_Into_Jsonl(512)
    # Merge_Emsplus_Into_Jsonl(1024)
    Merge_Emsplus_Into_Jsonl(333)

    # for SelectDataSet in range(len(DataSet)):
    #     if SelectDataSet == 1:
    #         continue
    #     Gen_Ext_Slice_Json(DataSet[SelectDataSet],256)

